package com.lw.leetcode.tree.b;

import java.util.*;

/**
 * Created with IntelliJ IDEA.
 * 6262. 图中最大星和
 *
 * @author liw
 * @version 1.0
 * @date 2022/12/11 12:34
 */
public class MaxStarSum {

    public int maxStarSum(int[] vals, int[][] edges, int k) {
        Map<Integer, List<Integer>> map = new HashMap<>();
        for (int[] edge : edges) {
            int a = edge[0];
            int b = edge[1];
            map.computeIfAbsent(a, v -> new ArrayList<>()).add(b);
            map.computeIfAbsent(b, v -> new ArrayList<>()).add(a);
        }
        int max = Integer.MIN_VALUE;
        int length = vals.length;
        PriorityQueue<Integer> queue = new PriorityQueue<>((a, b) -> Integer.compare(b, a));
        for (int i = 0; i < length; i++) {
            List<Integer> list = map.get(i);
            int v = vals[i];
            if (list == null) {
                max = Math.max(max, v);
                continue;
            }
            for (Integer index : list) {
                queue.add(vals[index]);
            }
            int j = 0;
            while (j < k && !queue.isEmpty()) {
                Integer p = queue.poll();
                if (p >0) {
                    v += p;
                } else {
                    break;
                }
                j++;
            }
            queue.clear();
            max = Math.max(max, v);
        }
        return max;
    }

}
